/****************************************************************************
*					LoopyBP
*
*	Description:	Runs inference for a given PGM using loopy belief propagation
*
****************************************************************************/

#ifndef _LoopyBP_H
#define _LoopyBP_H

#include <vector>
#include <map>
#include <set>
#include "PGMStruct.h"
#include "DataSet.h"
#include "Inference.h"
#include "ClusterGraph.h"

class LoopyBP : public Inference {
  private:

  protected:

	/* Data */
	  // Cluster graph for running belief propagation
	ClusterGraph clusterGraph;

	  // bool variable that shows that the last run of unconditional inference was successful (did converged)
	bool lastRunConvergence;

	  // damping parameter (used to get loopy belief propagation to converge sooner)
	float lambda;

	  // acceptable level of average squared error between belief sums over one separation set (defines convergence)
	float acceptableDeviation;

	  // Internal vector of sepSetIDs that contain sequince in which to send messages (needed for propagation)
	std::vector <unsigned int> sepSetSeq;
	  // Internal vector of bool values that show which sepSets were free in the sepSetSeq (needed for propagation)
	std::vector <bool> isSepSetFree;
	  // Internal vector of bool values that show which domain is left (or was visited in this path)
	std::vector <bool> isDomainFree;

	  // las value of Maximum Average Square Error:
	float lastMASE;
	  // number of iterations lastMASE holds
	unsigned int lastMASEHolds;



	/* Functions */
	  // Picks a sepSet from a vector of sepSets in clusterGraph
	bool selectSepSet (unsigned int& chosenVal, PGMStruct &pgmStruct);

	  // A recursive function that propagate beliefs along a tree that is constructed on the way
	void runPropagation (unsigned int& index, const bool& dir, PGMStruct& pgmStruct);

	  // Running message propagation along the constructed way in order to get expectation of the feature w.r.t. distribution defined by PGM
	float runCondPropagation (unsigned int& oneVarDom, std::vector < std::pair<unsigned int, bool> >& path, const std::vector < std::pair<unsigned int, unsigned int> >& featureDescr, const unsigned int& featureDescrSize, PGMStruct& pgmStruct);

	  // Returns the shortest path in cluster graph that covers fully the feature of interest
	int getMinPath (std::vector <unsigned int>& firstDomains, std::vector < std::vector < std::pair<unsigned int, bool> > >& path, const std::vector < std::pair<unsigned int, unsigned int> >& featureDescr, const unsigned int& featureDescrSize, PGMStruct& pgmStruct);
	int getMinPath (unsigned int& firstDomain, std::vector < std::pair<unsigned int, bool> >& path, std::vector < std::vector < std::pair<unsigned int, unsigned int> > >& subFeatureDescr, const std::vector < std::pair<unsigned int, unsigned int> >& featureDescr, const unsigned int& featureDescrSize, PGMStruct& pgmStruct);

	  // Creates path from root to the desired list domains
	void getPathFromRoot (std::vector <std::pair<unsigned int, bool>>& curPath, unsigned int& curLen, const unsigned int& curDomID, std::map< unsigned int, std::pair <std::vector<bool>, std::vector<unsigned int> > >& colors, std::vector<bool>& unCovered) const;

	  // Function sets all marginals to be from 0 to 1 (divides by partition functions) and sets a partition function value
	float setZandNormalizeBeliefs (void);

	  // recursive function that normalize belief for domain curDomID and all its sepSets and runs itself for all its descendatns
	float normalizeBelief (const unsigned int& curDomID, std::vector<bool>& domainCovered);

	  // An inline implementation of belief summation on all its variables values except for variable 'varToLeave'
	inline void sumBelief (const unsigned int& domainID, const unsigned int& varToLeave, PGMStruct &pgmStruct, std::vector <float>& message) const {
		sumBelief (clusterGraph.domainDescriptions[domainID].belief, domainID, varToLeave, pgmStruct, message);
	}

	  // An inline implementation of belief summation on all its variables values except for variable 'varToLeave' (with specification where to get current belief vector)
	inline void sumBelief (const std::vector<float>& curBelief, const unsigned int& domainID, const unsigned int& varToLeave, PGMStruct &pgmStruct, std::vector <float>& message) const {
		std::vector <unsigned int> vals;
		int u;
		int uMax = (int) pgmStruct.features[domainID].varsIDs.size();
		vals.resize(uMax, 0);
		// Prepearing message variable
		message.resize(pgmStruct.numOfValues[varToLeave], 0);
		for (unsigned int i = 0; i < message.size(); i++) 
			message[i] = 0;
		// getting index of variable 'varToLeave' in the domain pgmStruct.features[domainID]
		int varToLeaveIndex = 0;
		while ( (varToLeaveIndex < uMax) && (pgmStruct.features[domainID].varsIDs[varToLeaveIndex] != varToLeave) )
			varToLeaveIndex++;
		// setting uMax to point on the last value of 'vals'
		uMax--;
		// going through all domain variable values calulating a sum of belief
		unsigned int b = 0;
		do {
			// coiunting values 'vals' of the belief
			message[vals[varToLeaveIndex]] += curBelief[b];
			b++;
			// moving to the next set of values
			u = uMax;
			ClusterGraph::moveValues (vals, u, pgmStruct.features[domainID].varsIDs, pgmStruct.numOfValues);
		} while (u >= 0);
	}

	  // combining belief sum, old message sent in this direction and sepSet parameters (mu) into one vector (to use in belief update)
	inline void combinedMsg (std::vector <float>& resMsg, std::vector <float>& beliefSum, std::vector <float>& oldMsg, std::vector <float>& mu) const {
		for (unsigned int i = 0; i < resMsg.size(); i++)
			resMsg[i] = (lambda * beliefSum[i] + ((float) 1.0 - lambda)*oldMsg[i]) / mu[i];
	}

	  // (No Dumping version) combining belief sum, old message sent in this direction and sepSet parameters (mu) into one vector (to use in belief update)
	inline void combinedMsg (std::vector <float>& resMsg, std::vector <float>& beliefSum, std::vector <float>& mu) const {
		for (unsigned int i = 0; i < resMsg.size(); i++)
			resMsg[i] = beliefSum[i] / mu[i];
	}

	  // updates belief of a given domain given fully computed update message (including division of the message by the corresponding sepSet mu)
	inline void updateBelief (const unsigned int& domainID, const unsigned int& commonVar, std::vector <float>& updateMsg, PGMStruct &pgmStruct) {
		updateBelief (clusterGraph.domainDescriptions[domainID].belief, domainID, commonVar, updateMsg, pgmStruct);
	}

	  // updates belief given fully computed update message (including division of the message by the corresponding sepSet mu)
	inline void updateBelief (std::vector<float>& recBelief, const unsigned int& domainID, const unsigned int& commonVar, std::vector <float>& updateMsg, PGMStruct &pgmStruct) {
		std::vector <unsigned int> vals;
		int u;
		int uMax = (int) pgmStruct.features[domainID].varsIDs.size();
		vals.resize(uMax, 0);
		int commonVarIndex = 0;
		while ( (commonVarIndex < uMax) && (pgmStruct.features[domainID].varsIDs[commonVarIndex] != commonVar) )
			commonVarIndex++;
		uMax--;
		// going through all domain variable values and updating receiver belief 
		unsigned int b = 0;
		do {
			// updating belief
			recBelief[b] *= updateMsg[ vals[commonVarIndex] ];
			b++;
			// moving to the next set of values
			u = uMax;
			ClusterGraph::moveValues (vals, u, pgmStruct.features[domainID].varsIDs, pgmStruct.numOfValues);
		} while (u >= 0);
	}

	  // updates sep set parameter given its old value and fully computed update to the corresponding belief
	inline void updateSepSet (std::vector<float>& newMu, std::vector<float>& msg, std::vector<float>& oldMu) const {
		for (unsigned int i = 0; i < newMu.size(); i++) 
			newMu[i] = oldMu[i] * msg[i];
	}

	  // function for sending a message along sep set with sepSetID in direction 'dir'
	inline void sendMessage (const unsigned int& sepSetID, const bool& dir, PGMStruct &pgmStruct) {
		std::vector <float> message;
		if (dir) {// forward direction: from parOne to parTwo = from domain with one variable to domain with many variables
			// so, message is equal to belief of the sending domain
			// message = clusterGraph.domainDescriptions[clusterGraph.sepSets[sepSetID].parOne].belief;
			message.resize(clusterGraph.sepSets[sepSetID].muOneTwo.size());
			normalizeMsgToAveOne (message, clusterGraph.domainDescriptions[clusterGraph.sepSets[sepSetID].parOne].belief);
			if (clusterGraph.sepSets[sepSetID].direction) {
				combinedMsg (message, message, clusterGraph.sepSets[sepSetID].muOneTwo, clusterGraph.sepSets[sepSetID].muOneTwo);
				updateSepSet (clusterGraph.sepSets[sepSetID].muOneTwo, message, clusterGraph.sepSets[sepSetID].muOneTwo);
			} else {
				combinedMsg (message, message, clusterGraph.sepSets[sepSetID].muOneTwo, clusterGraph.sepSets[sepSetID].muTwoOne);
				updateSepSet (clusterGraph.sepSets[sepSetID].muOneTwo, message, clusterGraph.sepSets[sepSetID].muTwoOne);
			}
			// updating belief
			updateBelief (clusterGraph.sepSets[sepSetID].parTwo, clusterGraph.sepSets[sepSetID].commonVar, message, pgmStruct);
		} else {
			sumBelief (clusterGraph.sepSets[sepSetID].parTwo, clusterGraph.sepSets[sepSetID].commonVar, pgmStruct, message);
			normalizeMsgToAveOne (message, message);
			if (clusterGraph.sepSets[sepSetID].direction) {
				combinedMsg (message, message, clusterGraph.sepSets[sepSetID].muTwoOne, clusterGraph.sepSets[sepSetID].muOneTwo);
				updateSepSet (clusterGraph.sepSets[sepSetID].muTwoOne, message, clusterGraph.sepSets[sepSetID].muOneTwo);
			} else {
				combinedMsg (message, message, clusterGraph.sepSets[sepSetID].muTwoOne, clusterGraph.sepSets[sepSetID].muTwoOne);
				updateSepSet (clusterGraph.sepSets[sepSetID].muTwoOne, message, clusterGraph.sepSets[sepSetID].muTwoOne);
			}
			// updating belief
			updateBelief (clusterGraph.sepSets[sepSetID].parOne, clusterGraph.sepSets[sepSetID].commonVar, message, pgmStruct);
		}
		// updating sepSet direction
		clusterGraph.sepSets[sepSetID].direction = dir;
	}
	  // function for sending a message with NO DUMPING along sep set with sepSetID in direction 'dir' and using belief and saving result to the specified location
	inline void sendMessageConst (std::vector<float>& receiverBelief, std::vector<float>& senderBelief, const unsigned int& sepSetID, const bool& dir, PGMStruct &pgmStruct) {
		std::vector <float> message;
		if (dir) {// forward direction: from parOne to parTwo = from domain with one variable to domain with many variables
			// so, message is equal to belief of the sending domain
			//message.resize(clusterGraph.sepSets[sepSetID].muOneTwo.size());
			//normalizeMsgToAveOne (message, senderBelief);
			message = senderBelief;
			if (clusterGraph.sepSets[sepSetID].direction) {
				// old message is equal to new message (no dumping in case of tree propagation)
				combinedMsg (message, message, clusterGraph.sepSets[sepSetID].muOneTwo);
			} else {
				combinedMsg (message, message, clusterGraph.sepSets[sepSetID].muTwoOne);
			}
			// updating belief
			updateBelief (receiverBelief, clusterGraph.sepSets[sepSetID].parTwo, clusterGraph.sepSets[sepSetID].commonVar, message, pgmStruct);
		} else {
			sumBelief (senderBelief, clusterGraph.sepSets[sepSetID].parTwo, clusterGraph.sepSets[sepSetID].commonVar, pgmStruct, message);
			//normalizeMsgToAveOne (message, message);
			if (clusterGraph.sepSets[sepSetID].direction) {
				// old message is equal to new message (no dumping in case of tree propagation)
				combinedMsg (message, message, clusterGraph.sepSets[sepSetID].muOneTwo);
			} else {
				combinedMsg (message, message, clusterGraph.sepSets[sepSetID].muTwoOne);
			}
			// updating belief
			updateBelief (receiverBelief, clusterGraph.sepSets[sepSetID].parOne, clusterGraph.sepSets[sepSetID].commonVar, message, pgmStruct);
		}
	}

	  // normalizes message in a way that each element is now has a mean of 1
	inline void normalizeMsgToAveOne (std::vector<float>& messageNorm, std::vector<float>& messageOrig) const {
		float curSum = 0.0;
		for (unsigned int i = 0; i < messageOrig.size(); i++) 
			curSum += messageOrig[i];
		curSum /= (float) messageOrig.size();
		for (unsigned int i = 0; i < messageNorm.size(); i++) 
			messageNorm[i] = messageOrig[i] / curSum;
	}

  public:

	  // Standard constructor
	LoopyBP(void);
      
	  // Standard destructor
    ~LoopyBP();

	  // Sets parameters from environment (defined by config-file)
	int setParameters (Environment &environment);

      // Runs inference
	int run (float& logZ, PGMStruct& pgmStruct);

	  // Runs inference conditional on a feature 'featureDescr'
	int run (float& expectedValue, const std::vector < std::pair<unsigned int, unsigned int> >& featureDescr, PGMStruct& pgmStruct);

	  // Runs inference conditional on a feature 'featureDescr' with a given featureDescrSize
	int run (float& expectedValue, const std::vector < std::pair<unsigned int, unsigned int> >& featureDescr, const unsigned int& featureDescrSize, PGMStruct& pgmStruct);

	  // Runs inference for evaluation of expected value of a diven feature
	int run (float& expectedValue, Feature& feature, std::vector <unsigned int>& varIDs, PGMStruct& pgmStruct);

	  // Return likelihood of the data in the data set
	float getLogLikelihood (DataSet& dataSet, PGMStruct& pgmStruct);
	
}; // end of class

// Additional struct definition - it's used by internal functions of the LoopyBP
struct Triplet {
	std::vector<bool> reqVars;
	unsigned int sepSetID;
	bool dir;
};

#endif // _LoopyBP_H
